package cn.kinoko.aspect;

import cn.kinoko.annotation.RateLimiter;
import cn.kinoko.error.TryAcquireTokenFailException;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;

/**
 * 分布式锁切面
 *
 * @author kinoko
 */
@Aspect
@Component
public class RateLimiterAspect {

    @Autowired
    private RedissonClient redissonClient;

    @Around("@annotation(rateLimiter)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) throws Throwable {
        // 获取redisKey
        String lockName = this.getRedisKey(joinPoint, rateLimiter);
        Object target;
        // 获取分布式限流器
        RRateLimiter limiter = redissonClient.getRateLimiter(lockName);
        limiter.trySetRate(rateLimiter.rateType(), rateLimiter.rate(), rateLimiter.rateInterval(), rateLimiter.intervalUnit());
        // 尝试获取令牌
        boolean success = rateLimiter.timeout() == -1 ?
                limiter.tryAcquire(1) :
                limiter.tryAcquire(rateLimiter.timeout(), rateLimiter.timeUnit());
        // 令牌获取失败返回msg
        if (!success) {
            throw new TryAcquireTokenFailException(rateLimiter.errorDesc());
        }
        // 执行业务方法
        target = joinPoint.proceed();
        return target;
    }


    /**
     * 获取加锁的key
     *
     * @param joinPoint   切点
     * @param rateLimiter 注解
     * @return key
     */
    private String getRedisKey(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) {
        String key = rateLimiter.key();
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        // 获取切面方法
        Method method = signature.getMethod();
        // 如果key为空，则使用方法名作为key
        if (key != null && !key.isEmpty()) {
            key = method.getName();
        }
        // 设置限流粒度
        switch (rateLimiter.granularity()) {
            case USER:
                key += ":" + "uid";
                break;
            case IP:
                key += ":" + "ip";
                break;
            default:
        }
        return key;
    }

}